package main

import (
	"bytes"
	"elasticmigrate/utils"
	"flag"
	"fmt"
	jsoniter "github.com/json-iterator/go"
	"github.com/spf13/cast"
	"github.com/tidwall/gjson"
	"log/slog"
	"os"
	"strings"
	"sync"
	"time"
)

var (
	sourceAddress  string
	sourceUser     string
	sourcePassword string
	sourceIndex    string
	queryFile      string
	thread         int
	targetAddress  string
	targetUser     string
	targetPassword string
	targetIndex    string
	size           int
)

// 要过滤的字段列表
var filterFields []string

// 等待
var group = sync.WaitGroup{}

// 控制并发
var myChan chan struct{}

// json 第三方库
var json = jsoniter.ConfigCompatibleWithStandardLibrary

var query []byte

var bulkAction = "index"

// 初始化操作
func init() {
	// 命令行参数
	flag.StringVar(&sourceAddress, "sa", "http://localhost:9200", "源地址(sourceAddress), 默认为 http://localhost:9200")
	flag.StringVar(&sourceUser, "su", "elastic", "源用户(sourceUser, 默认为 elastic")
	flag.StringVar(&sourcePassword, "sp", "", "源密码(sourcePassword)")
	flag.StringVar(&queryFile, "q", "query.json", "查询条件query(参数为json文件路径), 默认为 query.json")
	flag.StringVar(&sourceIndex, "si", "", "源索引，可使用星号：keywords*")
	flag.StringVar(&targetAddress, "ta", "http://localhost:9200", "目标地址(targetAddress)")
	flag.StringVar(&targetUser, "tu", "elastic", "目标用户(targetUser), 默认为 elastic")
	flag.StringVar(&targetPassword, "tp", "", "目标密码(targetPassword)")
	flag.StringVar(&targetIndex, "ti", "", "目标索引")
	flag.StringVar(&bulkAction, "ba", "index", "bulk类型, 默认为index")
	flag.IntVar(&thread, "t", 20, "线程数(默认20)")
	flag.IntVar(&size, "n", 3000, "每次读取数量3000条")
	var field string
	flag.StringVar(&field, "ff", "", "目标密码(targetPassword)")
	flag.Parse()

	// 字段验证
	if sourceIndex == "" {
		slog.Error("源索引不能为空!")
		os.Exit(0)
	}
	if sourceAddress != "" && !strings.HasPrefix(sourceAddress, "http") {
		slog.Error("源服务器地址有误!")
		os.Exit(0)
	}
	if targetAddress != "" && !strings.HasPrefix(targetAddress, "http") {
		slog.Error("源服务器地址有误!")
		os.Exit(0)
	}
	if queryFile == "" {
		queryFile = "query.json"
	}
	if bulkAction != "" && bulkAction != "index" {
		bulkAction = "create"
	}
	fmt.Println(bulkAction)
	query, _ = os.ReadFile(queryFile)
	//if size < 500 || size > 10000 {
	//	slog.Warn("n不合符预期，自动设置为3000!")
	//}
	// 字段拆分
	filterFields = strings.Split(field, ",")
	myChan = make(chan struct{}, thread)
	for len(myChan) != cap(myChan) {
		myChan <- struct{}{}
	}
}

func main() {
	start := time.Now()
	// 验证两个服务器是否能正常访问
	sourceEsClient := utils.NewEsClient(sourceAddress, sourceUser, sourcePassword, sourceIndex, size, query)
	if !sourceEsClient.Ping("Source") {
		os.Exit(1)
	}
	count := sourceEsClient.GetCount()
	if count == 0 {
		slog.Error("源索引文档数为0 或 索引不存在!")
		return
	}
	slog.Info(fmt.Sprintf("Index: %s, Count: %d", sourceIndex, count))
	targetEsClient := utils.NewEsClient(targetAddress, targetUser, targetPassword, targetIndex, size, query)
	if !targetEsClient.Ping("Target") {
		os.Exit(1)
	}
	// 第一次查询 获取 滚动ID
	scrollId, data := sourceEsClient.FirstGetSourceIdAndData()
	group.Add(1)
	go func() {
		defer func() {
			group.Done()
		}()
		taskFunc(data, targetEsClient)
	}()
	for i := 0; i < count; i += size {
		scrollId, data = sourceEsClient.GetSourceIdAndData(scrollId)
		<-myChan
		group.Add(1)
		go func() {
			defer func() {
				myChan <- struct{}{}
				group.Done()
			}()
			taskFunc(data, targetEsClient)
		}()
	}
	group.Wait()
	slog.Info(fmt.Sprintf("程序运行完成，共耗时：%.2fs", float64(time.Now().Sub(start)/time.Millisecond)/1000))
}
func taskFunc(data []*utils.EsDoc, targetEsClient *utils.EsClient) {
	var bulkBuf = new(bytes.Buffer)
	for _, item := range data {
		// 过滤字段
		doc := utils.DeleteMapFields(filterFields, item.Source)
		var item1 = map[string]map[string]string{
			bulkAction: {
				"_index": item.Index,
			},
		}
		if bulkAction == "create" {
			item1[bulkAction]["_id"] = item.Id
		}
		if targetIndex != "" {
			// 设置索引
			var docIndex = targetIndex
			if strings.Contains(docIndex, "#y") ||
				strings.Contains(docIndex, "#m") ||
				strings.Contains(docIndex, "#d") {
				var date = time.Now()
				var err error
				if v, ok := doc["@timestamp"].(string); ok {
					date, err = cast.StringToDate(v)
					if err != nil {
						date = time.Now()
					}
				}
				docIndex = strings.ReplaceAll(docIndex, "#y", fmt.Sprintf("%d", date.Year()))
				docIndex = strings.ReplaceAll(docIndex, "#m", fmt.Sprintf("%d", date.Month()))
				docIndex = strings.ReplaceAll(docIndex, "#d", fmt.Sprintf("%d", date.Day()))
			}
			item1[bulkAction]["_index"] = docIndex
		}
		itemByte1, _ := json.Marshal(item1)
		bulkBuf.Write(itemByte1)
		bulkBuf.WriteString("\n")
		itemByte2, _ := json.Marshal(doc)
		bulkBuf.Write(itemByte2)
		bulkBuf.WriteString("\n")
	}
	// 关闭bulk添加
	res, _ := targetEsClient.EsBulkAction(bulkBuf.String())
	successNum := 0
	errorNum := 0
	gjson.Get(res.String(), "items").ForEach(func(key, value gjson.Result) bool {
		status := value.Get(bulkAction + ".status").Int()
		if status == 201 {
			// 迁移成功
			successNum++
		} else {
			// 迁移失败
			errorNum++
		}
		// 是否继续循环
		return true
	})
	slog.Info(fmt.Sprintf("[Dump] SuccessNum: %d, ErrorNum: %d", successNum, errorNum))
}
